#include "eye_custom_tiling.h"
#include "register/op_def_registry.h"
#include "graph/utils/type_utils.h"
#include "tiling/platform/platform_ascendc.h"

namespace optiling {
    // 通用计算对齐长度
    // static inline uint64_t AlignUp(uint64_t x, uint64_t a) { return (x + a - 1) / a * a; }
    // int32_t ld_aligned = (((int)d + 32/(int)typeSize - 1) / (32/(int)typeSize)) * (32/(int)typeSize);

    const uint32_t BLOCK_SIZE = 32;
    const uint32_t CORE_NUM = 8;
    const uint32_t BUFFER_NUM = 2;
    int32_t tilingKey = 0;

    static ge::graphStatus TilingFunc(gert::TilingContext* context)
    {
        EyeCustomTilingData tiling;
        // 动态获取硬件ub大小和核心数量
        uint64_t ubSize;
        auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
        ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
        auto coreNum = ascendcPlatform.GetCoreNum();
        // 获取输入张量形状和数据类型
        uint32_t inputNum = context->GetInputShape(0)->GetStorageShape().GetShapeSize();        // m*n
        uint32_t typeLength = 0;
        ge::TypeUtils::GetDataTypeLength(context->GetInputDesc(0)->GetDataType(), typeLength);  // float32->4
        uint32_t inputLength = inputNum * typeLength;                                           // m*n*4
        uint32_t typeSize = inputLength / inputNum;                                             // 4 inputBytes->typeSize
        uint32_t alignNum = BLOCK_SIZE / typeSize; // 每行对齐到32字节，计算对齐后的元素数量
        
        const auto Shape = context->GetInputTensor(0)->GetOriginShape();
        uint32_t rowLength     = Shape.GetDim(0);
        uint32_t columnLength  = Shape.GetDim(1);
        uint32_t diagLen = (rowLength < columnLength) ? rowLength : columnLength;
        uint32_t indexK  = 0; 
        // 核间-计算每个核处理的对角线长度
        uint32_t tailBlockLength = diagLen / CORE_NUM;
        uint32_t fullBlockNum = diagLen % CORE_NUM ;
        uint32_t tailBlockNum = CORE_NUM - fullBlockNum; // 对齐是全是尾核
        uint32_t fullBlockLength = tailBlockLength + 1;

        // TilingKey 编码规则:
            // input
            // half        -> 0   
            // float       -> 1
            // int16_t     -> 2
            // int32_t     -> 3
        const auto inputDataType = context->GetInputDesc(0)->GetDataType();

        switch (inputDataType) {
            case ge::DT_FLOAT16: tilingKey = 0; break; 
            case ge::DT_FLOAT:   tilingKey = 1; break; 
            case ge::DT_INT16:   tilingKey = 2; break; 
            case ge::DT_INT32:   tilingKey = 3; break; 

            default: return ge::GRAPH_FAILED;
        }
        context->SetTilingKey(tilingKey);

        context->SetBlockDim(CORE_NUM);

        tiling.set_typeSize(typeSize);
        tiling.set_alignNum(alignNum);
        tiling.set_matrixOrder(0); 
        
        tiling.set_rowLength(rowLength);
        tiling.set_columnLength(columnLength);
        tiling.set_indexK(indexK);
        tiling.set_diagLen(diagLen);
        tiling.set_fullBlockLength(fullBlockLength);
        tiling.set_tailBlockLength(tailBlockLength);
        tiling.set_fullBlockNum(fullBlockNum);
        tiling.set_tailBlockNum(tailBlockNum);
        
        tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
        context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
        size_t *currentWorkspace = context->GetWorkspaceSizes(1);
        currentWorkspace[0] = 0;
        return ge::GRAPH_SUCCESS;
    }
}

namespace ge {
    static ge::graphStatus InferShape(gert::InferShapeContext* context)
    {
        const gert::Shape* x1_shape = context->GetInputShape(0);
        gert::Shape* y_shape = context->GetOutputShape(0);
        *y_shape = *x1_shape;
        return GRAPH_SUCCESS;
    }
    static ge::graphStatus InferDataType(gert::InferDataTypeContext *context)
    {
        const auto inputDataType = context->GetInputDataType(0);
        context->SetOutputDataType(0, inputDataType);
        return ge::GRAPH_SUCCESS;
    }
}

namespace ops {
class EyeCustom : public OpDef {
public:
    explicit EyeCustom(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT16, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT16, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType);
        this->AICore().SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");
    }
};

OP_ADD(EyeCustom);
}
